#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
四足机器人模块
包含四足机器人类，负责机器人模型和运动学计算
"""

import numpy as np
import ikpy.chain
import os
import pybullet as p
import pybullet_data

# 导入配置
from config import ROBOT_START_POS, ROBOT_START_ORI, USE_GUI

class FourLeggedRobot:
    """四足机器人类，管理机器人模型和运动学"""
    
    def __init__(self, urdf_path):
        """初始化四足机器人
        
        参数:
            urdf_path: URDF文件路径
        """
        self.urdf_path = urdf_path
        self.robot_id = None
        self.joint_name_to_id = {}
        self.leg_chains = {}
        self.base_link_world_pos_np = None
        
        # 检查URDF文件是否存在
        if not os.path.exists(urdf_path):
            print(f"错误: URDF文件不存在: {urdf_path}")
            raise FileNotFoundError(f"错误: URDF文件不存在: {urdf_path}")
    
    def init_pybullet(self,yaw):
        """初始化PyBullet仿真环境
        
        返回:
            robot_id: 机器人ID
            joint_name_to_id: 关节名称到ID的映射
            base_link_world_pos_np: 基座位置
        """
        # 连接到PyBullet
        if USE_GUI:
            try:
                p.connect(p.GUI)
            except p.error:
                print("警告: 无法连接到GUI模式，回退到DIRECT模式")
                p.connect(p.DIRECT)
        else:
            p.connect(p.DIRECT)
            
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.8)
        
        # 加载地面
        plane_id = p.loadURDF('plane.urdf')
        p.changeDynamics(plane_id, -1, lateralFriction=1)
        
        # 加载机器人
        start_pos = ROBOT_START_POS
        start_ori = p.getQuaternionFromEuler(ROBOT_START_ORI)
        self.robot_id = p.loadURDF(self.urdf_path, start_pos, start_ori, useFixedBase=True)
        self.base_link_world_pos_np = np.array(start_pos)
        
        # 获取关节信息
        num_joints = p.getNumJoints(self.robot_id)
        print(f"机器人关节数量: {num_joints}")
        
        for i in range(num_joints):
            info = p.getJointInfo(self.robot_id, i)
            name = info[1].decode('utf-8')
            self.joint_name_to_id[name] = i
            print(f"{i}: {name}")
        
        # 设置仿真步长
        p.setTimeStep(1/240)
        
        # 设置摄像机
        if USE_GUI:
            p.resetDebugVisualizerCamera(
                cameraDistance=1.0,
                cameraYaw=yaw,
                cameraPitch=-30,
                cameraTargetPosition=[0, 0, 0.5]
            )
        
        return self.robot_id, self.joint_name_to_id, self.base_link_world_pos_np
    
    def create_leg_chains(self):
        """创建四条腿的运动链
        
        返回:
            leg_chains: 腿部运动链字典
        """
        try:
            # 左前腿 (LF)
            print("正在从URDF文件创建左前腿运动链...")
            self.leg_chains["LF"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "A1", "A1_Link", "A2", "A2_Link", "A3", "A3_Link", "A4"],
                active_links_mask=[False, True, True, True, False],
                name="left_front_leg"
            )
            
            # 右前腿 (RF)
            print("正在从URDF文件创建右前腿运动链...")
            self.leg_chains["RF"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "B1", "B1_Link", "B2", "B2_Link", "B3", "B3_Link", "B4"],
                active_links_mask=[False, True, True, True, False],
                name="right_front_leg"
            )
            
            # 左后腿 (LB)
            print("正在从URDF文件创建左后腿运动链...")
            self.leg_chains["LB"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "C1", "C1_Link", "C2", "C2_Link", "C3", "C3_Link", "C4"],
                active_links_mask=[False, True, True, True, False],
                name="left_back_leg"
            )
            
            # 右后腿 (RB)
            print("正在从URDF文件创建右后腿运动链...")
            self.leg_chains["RB"] = ikpy.chain.Chain.from_urdf_file(
                self.urdf_path,
                base_elements=["base_link", "D1", "D1_Link", "D2", "D2_Link", "D3", "D3_Link", "D4"],
                active_links_mask=[False, True, True, True, False],
                name="right_back_leg"
            )
            
            return self.leg_chains
            
        except Exception as e:
            print(f"创建运动链时出错: {e}")
            raise
    
    def reset_leg_positions(self):
        """重置所有腿到初始位置"""
        for leg_id, chain in self.leg_chains.items():
            zero_joint_angles = np.zeros(len(chain.links))
            self.apply_joint_angles(chain, zero_joint_angles)
    
    def apply_joint_angles(self, chain, angles):
        """应用关节角度到机器人

        参数:
            chain: 运动链
            angles: 关节角度数组
        """
        for i, mask in enumerate(chain.active_links_mask):
            if mask:
                name = chain.links[i].name
                if name in self.joint_name_to_id:
                    p.setJointMotorControl2(
                        self.robot_id,
                        self.joint_name_to_id[name],
                        p.POSITION_CONTROL,
                        targetPosition=angles[i],
                        force=500
                    )
    
    def solve_ik_and_apply(self, leg_id, target_pos_base_link, init_pos=None):
        """IK求解并应用关节角度
        
        参数:
            leg_id: 腿的ID (LF, RF, LB, RB)
            target_pos_base_link: 目标位置 (基座坐标系)
            init_pos: 初始猜测位置
            
        返回:
            solution: 关节角度解
            error: 位置误差
        """
        chain = self.leg_chains[leg_id]
        
        # 如果没有提供初始猜测，使用零向量
        if init_pos is None:
            init_pos = np.zeros(len(chain.links))
        
        # 求解IK
        solution = chain.inverse_kinematics(
            target_pos_base_link,
            initial_position=init_pos,
            max_iter=1000
        )
        
        # 应用解到机器人
        self.apply_joint_angles(chain, solution)
        
        # 计算误差
        fk = chain.forward_kinematics(solution)
        pos_base_link = fk[:3, 3]
        error = np.linalg.norm(pos_base_link - target_pos_base_link)
        return solution, error
    
    def get_world_position(self, base_link_position):
        """将基座坐标系位置转换为世界坐标系位置
        参数:base_link_position: 基座坐标系位置   
        返回:world_position: 世界坐标系位置
        """
        return (base_link_position + self.base_link_world_pos_np).tolist()
